'''
完成策略的回测，绘制以沪深300为基准的收益曲线，计算年化收益、最大回撤、夏普比率
主要的方法包括:
isKUpBreakMA10: 当日K线是否上穿10日均线
isKDownBreakMA10: 当日K线是否下穿10日均线
compareClost2MA10: 工具方法，某日收盘价和当日对应的10均线的关系
backTest: 回测主逻辑方法， 从股票池获取股票后，按照每天的交易日一天天回测
computeDrawDown: 计算最大回撤
computeAnnualProfit: 计算年化收益
computeSharpeRatio: 计算复普比率
'''

import matplotlib.pyplot as plt
import pandas as pd
from pymongo import MongoClient, ASCENDING, DESCENDING

from pe_ttm.StockPoolStrategy import stockPool, findOutStocks

db = MongoClient('mongodb://localhost:27017')['quant']

'''
判断某只股票在某日是否满足K线上穿10日均线
:param code: 股票代码
:param date: 日期
:return: True/False
'''
def isKUpBreakMA10(code, date):
    # 如果股票当日停牌或者是下跌，则返回False
    currentDaily = db.daily_k.find_one(
        {'code': code, 'time': date, 'trade_status': '交易'}
    )

    if currentDaily is None or currentDaily['pct_chg'] < 0:
        return False

    # 计算MA10
    dailyCursor = db.daily_k.fin(
        {'code': code, 'time': {'$lte': date}},
        sort = [('time', DESCENDING)],
        limit = 11,
        projection = {'close': True, 'adjfactor': True, 'trade_status': True}
    )

    dailies = [x for x in dailyCursor]

    if len(dailies) < 11:
        return False

    dailies.reverse()

    lastClose2LastMA10 = compareClose2MA10(dailies[0:10])
    currentClose2CurrentMA10 = compareClose2MA10(dailies[1:])

    if lastClose2LastMA10 is None or currentClose2CurrentMA10 is None:
        return False

    # 判断收盘价和MA10的大小
    return lastClose2LastMA10 <= 0 and currentClose2CurrentMA10 == 1

'''
判断某只股票在某日是否满足K线下穿10日均线

:param code: 股票代码
:param date: 日期
:return: True/False
'''
def isKDownBreakMA10(code, date):

    # 如果股票当日停牌或者是下跌，则返回False
    currentDaily = db.daily_k.find_one(
        {'code': code, 'time': date, 'trade_status': '交易'})
    if currentDaily is None or currentDaily['pct_chg'] > 0:
        return False

    # 计算MA10
    dailyCursor = db.daily_k.find(
        {'code': code, 'time': {'$lte': date}},
        sort = [('time', DESCENDING)],
        limit = 10,
        projection = {'close': True, 'adjfactor': True, 'trade_status': True}
    )

    dailies = [x for x in dailyCursor]

    if len(dailies) < 11:
        return False

    dailies.reverse()

    lastClose2LastMA10 = compareClose2MA10(dailies[0:10])
    currentClose2CurrentMA10 = compareClose2MA10(dailies[1:])

    if lastClose2LastMA10 is None or currentClose2CurrentMA10 is None:
        return False

    # 判断收盘价和MA10的大小
    return lastClose2LastMA10 >= 0 and currentClose2CurrentMA10 == -1

'''
策略回测。结束后打印出收益曲线(沪深300基准)、年化收益、最大回撤、

:param begin_date: 回测开始日期
:param end_date: 回测结束日期
'''
def backTest(beginDate, endDate):
    totalCapital = 1E7
    cash = 1E7
    totalValue = 0
    singlePosition = 2E5

    # 时间为key的净值，收益和同期沪深基准
    dfProfit = pd.DataFrame(columns=['net_value', 'profit', 'hs300'])

    allDatesCursor = db.daily_k.find(
        {'code': '000001.SH', 'time': {'$lte': endDate, '$gte': beginDate}},
        sort = [('time', ASCENDING)],
        projection = {'time':True},
        batch_size = 1000
    )

    allDates = [x['time'] for x in allDatesCursor]

    hs300BeginValue = db.daily_k.find_one(
        {'code': '000300.SH', 'time': allDates[0]},
        projection = {'close': True,}
    )['close']

    adjustDates, dateCodesDict = stockPool(beginDate, endDate)

    lastPhaseCodes = None
    thisPhaseCodes = None
    toBeSoldCodes = set()
    toBeBoughtCodes = set()
    holdingCodeDict = dict()
    lastDate = None

    # 按照日期一步步回测
    for date in allDates:
        print("BackTest at ".format(date))

        # 当期持仓股票列表
        beforeSellHoldingCodes = list(holdingCodeDict.keys())

        # 处理复权
        if lastDate is not None and len(beforeSellHoldingCodes) > 0:
            lastDailyCursor = db.daily_k.find(
                {'code': {'$in': beforeSellHoldingCodes}, 'time': lastDate},
                projection = {'code': True, 'adjfactor': True}
            )

            codeLastAdjactorDict = dict()
            for lastDaily in lastDailyCursor:
                codeLastAdjactorDict[lastDaily['code']] = lastDaily['adjfactor']

            currentDailyCursor = db.daily_k.find(
                {'code': {'$in': beforeSellHoldingCodes}, 'time': date},
                projection = {'code': True, 'adjfactor': True}
            )

            for currentDaily in currentDailyCursor:
                currentAdjfactor = currentDaily['adjfactor']
                code = currentDaily['code']
                beforeVolume = holdingCodeDict[code]['volume']
                lastAdjfactor = codeLastAdjactorDict[code]
                afterVolume = int(beforeVolume * (currentAdjfactor / lastAdjfactor))
                holdingCodeDict[code]['volume'] = afterVolume
                print('持仓量调整：{0},{1}, {2}, {3}, {4}'.format(code, beforeVolume, lastAdjfactor, afterVolume, currentAdjfactor))

        # 卖出
        if len(toBeSoldCodes) > 0:
            sellDailyCursor = db.daily_k.find(
                {'code': {'$in': list(toBeSoldCodes)}, 'time': date, 'trade_status': '交易'},
                projection = {'open': True, 'code': True}
            )

            for sellDaily in sellDailyCursor:
                code = sellDaily['code']
                if code in beforeSellHoldingCodes:
                    holdingStock = holdingCodeDict[code]
                    holdingVolume = holdingStock['volume']
                    sellPrice = sellDaily['open']
                    sellAmount = holdingVolume * sellPrice
                    cash += sellAmount

                    cost = holdingStock['cost']
                    singleProfit = (sellAmount - cost) * 100 / cost
                    print('卖出%s, %6d, %6.2f, %8.2f, %4.2f' %(code, holdingVolume, sellPrice, sellAmount, singleProfit))

                    del holdingCodeDict[code]
                    toBeSoldCodes.remove(code)
        print('卖出后，现金：{0}'.format(cash))

        # 买入
        if len(toBeBoughtCodes) > 0:
            buyDailyCursor = db.daily_k.find(
                {'code': {'$in', list(toBeBoughtCodes)}, 'time': date, 'trade_status': '交易'},
                projection = {'code': True, 'open': True}
            )

            for buyDaily in buyDailyCursor:
                if cash > singlePosition:
                    buyPrice = buyDaily['open']
                    code = buyDaily['code']
                    volume = int(int(singlePosition/buyPrice)/ 100) * 100
                    buyAmount = buyPrice * volume
                    cash -= buyAmount
                    holdingCodeDict[code] = {
                        'volume': volume,
                        'cost': buyAmount,
                        'last_value': buyAmount
                    }

                    print('买入%s, %6d, %6.2f, %8.2f' % (code, volume, buyPrice, buyAmount))

        print('买入后，现金: %10.2f' % cash)

        # 持仓股代码列表
        holdingCodes = list(holdingCodeDict.keys())

        # 如果调整日， 则获取新一期的股票列表
        if date in adjustDates:
            # 暂存为上期的日期
            if thisPhaseCodes is not None:
                lastPhaseCodes = thisPhaseCodes
            thisPhaseCodes = dateCodesDict[date]

            # 找到所有调出股票代码， 在第二日开盘时卖出
            if lastPhaseCodes is not None:
                outCodes = findOutStocks(lastPhaseCodes, thisPhaseCodes)
                for outCode in outCodes:
                    if outCode in holdingCodeDict:
                        toBeSoldCodes.add(outCode)

        # 检查是否有需要第二天卖出的股票
        for holdingCode in holdingCodes:
            if isKDownBreakMA10(holdingCode, date):
                toBeSoldCodes.add(holdingCode)

        # 检查是否有需要第二天买入的股票
        toBeBoughtCodes.clear()
        if thisPhaseCodes is not None:
            for code in thisPhaseCodes:
                if code not in holdingCode and isKUpBreakMA10(code, date):
                    toBeBoughtCodes.add(code)

        # 计算总资产
        totalValue = 0
        holdingDailyCursor = db.daily_k.find(
            {'code': {'$in': holdingCodes}, 'time': date},
            projection={'close': True, 'code': True}
        )
        for holdingDaily in holdingDailyCursor:
            code = holdingDaily['code']
            holdingStock = holdingCodeDict[code]
            value = holdingDaily['close'] * holdingStock['volume']
            totalValue += value

            profit = (value - holdingStock['cost']) * 100/holdingStock['cost']
            oneDayProfit = (value - holdingStock['last_value']) * 100 / holdingStock['last_value']

            holdingStock['last_value'] = value
            # 如果亏损大于2%，则止损
            if profit <= -2 or oneDayProfit<= -1:
                toBeSoldCodes.add(code)
            print('持仓: %s, %10.2f, %4.2f, %4.2f' %(code, value, profit, oneDayProfit))

        totalCapital = totalValue + cash
        hs300CurrentValue = db.daily_k.find_one(
            {'code': '000300.SH', 'time': date},
            projection = {'close': True}
        )['close']

        print('收盘后，现金: %10.2f, 总资产: %10.2f' % (cash, totalCapital))
        lastDate = date
        dfProfit.loc[date] = {
            'net_value': round(totalCapital/1e7, 2),
            'profit': round(100 * (totalCapital - 1e7)/1e7, 2),
            'hs300': round(100 * (hs300CurrentValue - hs300BeginValue)/ hs300BeginValue, 2)

        }

    print(dfProfit)
    drawDown = computeDrawDown(dfProfit['net_value'])
    annualProfit, sharpeRatio = computeSharpeRatio(dfProfit['net_value'])

    print('回测结果 %s - %s，年化收益： %7.3f, 最大回撤：%7.3f, 夏普比率：%4.2f' %(beginDate, endDate, annualProfit, drawDown, sharpeRatio))

    dfProfit.plot(title='BackTest Reusl', y = ['profit', 'hs300'], kind = 'line')
    plt.show()



'''
比较当前的收盘价和MA10的关系
:param dailies: 日线列表，10个元素，最后一个是当前交易日
:return 0 相等， 1 大于， -1小于， None结果未知
'''
def compareClose2MA10(dailies):
    currentDaily = dailies[9]
    closeSum = 0
    count = 0

    for daily in dailies:
        # 10天当中，只要有天停牌则返回False
        if daily['trade_status'] == '停牌' or 'adjfactor' not in daily:
            return None

        # 用后复权累计
        closeSum += daily['close'] * daily['adjfactor']
        count += 1

    if count < 10:
        return None

    # 计算MA10
    ma10 = closeSum / 10

    # 判断收盘价和MA10的大小
    postAdjustedClose = currentDaily['close'] * currentDaily['adjfactor']
    differ = postAdjustedClose - ma10

    if differ > 0:
        return 1
    elif differ < 0:
        return  -1
    else:
        return  0

'''
计算最大回撤
:param net_values: 净值列表
'''
def computeDrawDown(netValues):
    # 最大回撤的初始值为0
    maxDrawDown = 0
    size = len(netValues)
    index = 0

    # 双层循环找出最大回撤
    for netValue in netValues:
        for subNetValue in netValues[index:]:
            drawDown = 1 - subNetValue / netValue
            if drawDown > maxDrawDown:
                maxDrawDown = drawDown

        index += 1
    return maxDrawDown

'''
计算年化收益
'''
def computeAnnualProfit(tradingDays, netValue):
    annualProfit = 0

    if tradingDays > 0:
        # 计算年数
        years = tradingDays / 245
        # 计算年化收益
        annualProfit = pow(netValue, 1/years) - 1

    annualProfit = round(annualProfit * 100, 2)
    return  annualProfit

'''
计算夏普比率
:param net_values: 净值列表
'''
def computeSharpeRatio(netValues):
    # 总交易日数
    tradingDays = len(netValues)
    # 所有收益的DataFrame
    profitDF = pd.DataFrame(columns={'profit'})
    # 收益之后，初始化为第一天的收益
    profitDF.loc[0] = {'profit': round((netValues[0]-1) * 100, 2)}
    # 计算每天的收益
    for index in range(1, tradingDays):
        # 计算每日的收益变化
        profit = (netValues[index] - netValues[index - 1])/netValues[index - 1]
        profit = round(profit * 100, 2)
        profitDF.loc[index] = {'profit': profit}

    # 计算标准差
    profitSTD = pow(profitDF.var()['profit']/245, 1/2)

    # 年化收益
    annualProfit = computeAnnualProfit(tradingDays, netValues[-1])

    # 夏普比率
    sharpeRatio = (annualProfit - 4.75) / profitSTD

    return  annualProfit, sharpeRatio

if __name__ == '__main__':
    netValues = [1.6, 1.7, 0.6]
    print(computeDrawDown(netValues))

